from torch.utils.data import Dataset
import os
from PIL import Image
import random
from torchvision import transforms
from .pipelines import Indentity
from .utils import collect_data,eval_pipeline

AVERGAE_LENGTH = 10000

class BaseDataset(Dataset):
    def __init__(self,image_dir,gt_dir,mode='train',size=256,pipelines=[],pair=False,name='gan') -> None:
        super().__init__()
        
        assert mode in ['test','train']
        self.mode = mode
        self.size = size
        self.image_dir = image_dir
        self.gt_dir = gt_dir
        # to avoid the empty
        self.pipelines = [Indentity()]+[eval_pipeline(pipeline) for pipeline in pipelines]
        self.pair = pair
        self.name = name
        
        
        self.images = collect_data(self.image_dir)
        self.gts = collect_data(self.gt_dir)
        
        
        self.T = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
        ])
        self.flip = transforms.RandomHorizontalFlip(p=1)
        print('information:','\t','IMAGEs:','\t',len(self.images),'\t','GTs:','\t',len(self.gts))
    
    def __len__(self):
        return AVERGAE_LENGTH
    
    def __getitem__(self,index):
        if self.pair:
            # 根据gt去获取raw
            gt_style_path = random.choice(self.gts)
            image_path = gt_style_path.replace(self.gt_dir,self.image_dir)
        else:
            gt_style_path = random.choice(self.gts)
            image_path = random.choice(self.images)
        
        
        image = Image.open(image_path).convert('RGB')
        gt = Image.open(gt_style_path).convert('RGB')
        
        image = image.resize(size=(self.size,self.size))
        gt = gt.resize(size=(self.size,self.size))
        
        if self.mode == 'test':
            return self.T(image),self.T(gt)
        image,gt = self.make_aug(image,gt)
        
        return self.T(image),self.T(gt)
    
    def make_aug(self,image,gt):
        result = {}
        result['image'] = image
        result['gt'] = gt
        
        for pipeline in self.pipelines:
            result = pipeline(result)
        return result['image'],result['gt']
        
        
        
        